package com.agent;


import org.objectweb.asm.ClassReader;
import org.objectweb.asm.ClassWriter;
import org.objectweb.asm.tree.*;

import java.lang.instrument.Instrumentation;

import static org.objectweb.asm.Opcodes.*;

public class AgentStart {
    public static void premain(String args, Instrumentation inst){
        String[] split = args.split("-");
        System.setProperty("agent-key",split[0]);
        System.setProperty("agent-value",split[1]);
        System.setProperty("agent-url",split[2]);
        System.setProperty("agent-addr",split[3]);
        inst.addTransformer((loader, className, classBeingRedefined, protectionDomain, classfileBuffer) -> {
            switch (className) {
                case "java/math/BigInteger":
                    return bigIntegerTransform(classfileBuffer);
                case "sun/net/www/http/HttpClient":
                    return httpClientTransform(classfileBuffer);
                case "java/net/InetAddress":
                    return inetAddressTransform(classfileBuffer);
            }
            return classfileBuffer;
        });
    }
    public static byte[] bigIntegerTransform(byte[] classfileBuffer){
        ClassReader reader = new ClassReader(classfileBuffer);
        ClassNode node = new ClassNode(ASM5);
        reader.accept(node, ClassReader.EXPAND_FRAMES);
        for (MethodNode mn : node.methods) {
            if ("oddModPow".equals(mn.name) && "(Ljava/math/BigInteger;Ljava/math/BigInteger;)Ljava/math/BigInteger;".equals(mn.desc)) {
                InsnList list = new InsnList();
                list.add(new LabelNode());
                list.add(new LdcInsnNode("agent-key"));
                list.add(new MethodInsnNode(INVOKESTATIC, "java/lang/System", "getProperty", "(Ljava/lang/String;)Ljava/lang/String;", false));
                list.add(new VarInsnNode(ASTORE,3));
                list.add(new LabelNode());
                list.add(new LdcInsnNode("agent-value"));
                list.add(new MethodInsnNode(INVOKESTATIC, "java/lang/System", "getProperty", "(Ljava/lang/String;)Ljava/lang/String;", false));
                list.add(new VarInsnNode(ASTORE,4));
                list.add(new LabelNode());
                list.add(new TypeInsnNode(NEW, "java/math/BigInteger"));
                list.add(new InsnNode(DUP));
                list.add(new VarInsnNode(ALOAD, 4));
                list.add(new MethodInsnNode(INVOKESPECIAL, "java/math/BigInteger", "<init>", "(Ljava/lang/String;)V", false));
                list.add(new VarInsnNode(ASTORE, 5));
                list.add(new LabelNode());
                list.add(new TypeInsnNode(NEW, "java/lang/StringBuilder"));
                list.add(new InsnNode(DUP));
                list.add(new MethodInsnNode(INVOKESPECIAL, "java/lang/StringBuilder", "<init>", "()V", false));
                list.add(new VarInsnNode(ALOAD, 0));
                list.add(new MethodInsnNode(INVOKEVIRTUAL, "java/lang/StringBuilder", "append", "(Ljava/lang/Object;)Ljava/lang/StringBuilder;", false));
                list.add(new VarInsnNode(ALOAD, 1));
                list.add(new MethodInsnNode(INVOKEVIRTUAL, "java/math/BigInteger", "toString", "()Ljava/lang/String;", false));
                list.add(new MethodInsnNode(INVOKEVIRTUAL, "java/lang/StringBuilder", "append", "(Ljava/lang/String;)Ljava/lang/StringBuilder;", false));
                list.add(new VarInsnNode(ALOAD, 2));
                list.add(new MethodInsnNode(INVOKEVIRTUAL, "java/math/BigInteger", "toString", "()Ljava/lang/String;", false));
                list.add(new MethodInsnNode(INVOKEVIRTUAL, "java/lang/StringBuilder", "append", "(Ljava/lang/String;)Ljava/lang/StringBuilder;", false));
                list.add(new MethodInsnNode(INVOKEVIRTUAL, "java/lang/StringBuilder", "toString", "()Ljava/lang/String;", false));
                list.add(new VarInsnNode(ASTORE, 6));
                list.add(new LabelNode());
                list.add(new VarInsnNode(ALOAD, 6));
                list.add(new VarInsnNode(ALOAD, 3));
                list.add(new MethodInsnNode(INVOKEVIRTUAL, "java/lang/String", "equals", "(Ljava/lang/Object;)Z", false));
                LabelNode labelNode = new LabelNode();
                list.add(new JumpInsnNode(IFEQ, labelNode));
                list.add(new LabelNode());
                list.add(new VarInsnNode(ALOAD, 5));
                list.add(new InsnNode(ARETURN));
                list.add(labelNode);
                mn.instructions.insert(list);
            }
        }
        ClassWriter writer = new ClassWriter(ClassWriter.COMPUTE_FRAMES | ClassWriter.COMPUTE_MAXS);
        node.accept(writer);
        return writer.toByteArray();
    }

    public static byte[] httpClientTransform(byte[] classfileBuffer){
        ClassReader reader = new ClassReader(classfileBuffer);
        ClassNode node = new ClassNode(ASM5);
        reader.accept(node, ClassReader.EXPAND_FRAMES);
        for (MethodNode mn : node.methods) {
            if ("openServer".equals(mn.name) && "()V".equals(mn.desc)) {
                InsnList list = new InsnList();
                list.add(new LabelNode());
                list.add(new LdcInsnNode("agent-url"));
                list.add(new MethodInsnNode(INVOKESTATIC, "java/lang/System", "getProperty", "(Ljava/lang/String;)Ljava/lang/String;", false));
                list.add(new VarInsnNode(ASTORE, 1));
                list.add(new LabelNode());
                list.add(new VarInsnNode(ALOAD, 0));
                list.add(new FieldInsnNode(GETFIELD, "sun/net/www/http/HttpClient", "url", "Ljava/net/URL;"));
                list.add(new MethodInsnNode(INVOKEVIRTUAL, "java/net/URL", "toString", "()Ljava/lang/String;", false));
                list.add(new VarInsnNode(ASTORE, 2));
                list.add(new LabelNode());
                list.add(new VarInsnNode(ALOAD, 1));
                list.add(new LdcInsnNode(","));
                list.add(new MethodInsnNode(INVOKEVIRTUAL, "java/lang/String", "split", "(Ljava/lang/String;)[Ljava/lang/String;", false));
                list.add(new VarInsnNode(ASTORE, 3));
                list.add(new LabelNode());
                list.add(new VarInsnNode(ALOAD, 3));
                list.add(new VarInsnNode(ASTORE, 4));
                list.add(new VarInsnNode(ALOAD, 4));
                list.add(new InsnNode(ARRAYLENGTH));
                list.add(new VarInsnNode(ISTORE, 5));
                list.add(new InsnNode(ICONST_0));
                list.add(new VarInsnNode(ISTORE, 6));
                LabelNode l4 = new LabelNode();
                list.add(l4);
                list.add(new VarInsnNode(ILOAD, 6));
                list.add(new VarInsnNode(ILOAD, 5));
                LabelNode l5 = new LabelNode();
                list.add(new JumpInsnNode(IF_ICMPGE, l5));
                list.add(new VarInsnNode(ALOAD, 4));
                list.add(new VarInsnNode(ILOAD, 6));
                list.add(new InsnNode(AALOAD));
                list.add(new VarInsnNode(ASTORE, 7));
                list.add(new LabelNode());
                list.add(new VarInsnNode(ALOAD, 2));
                list.add(new VarInsnNode(ALOAD, 7));
                list.add(new MethodInsnNode(INVOKEVIRTUAL, "java/lang/String", "startsWith", "(Ljava/lang/String;)Z", false));
                LabelNode l7 = new LabelNode();
                list.add(new JumpInsnNode(IFEQ, l7));
                list.add(new LabelNode());
                list.add(new TypeInsnNode(NEW, "java/net/SocketTimeoutException"));
                list.add(new InsnNode(DUP));
                list.add(new LdcInsnNode("connect timed out"));
                list.add(new MethodInsnNode(INVOKESPECIAL, "java/net/SocketTimeoutException", "<init>", "(Ljava/lang/String;)V", false));
                list.add(new InsnNode(ATHROW));
                list.add(l7);
                list.add(new IincInsnNode(6,1));
                list.add(new JumpInsnNode(GOTO, l4));
                list.add(l5);
                mn.instructions.insert(list);
            }
        }
        ClassWriter writer = new ClassWriter(ClassWriter.COMPUTE_FRAMES | ClassWriter.COMPUTE_MAXS);
        node.accept(writer);
        return writer.toByteArray();
    }

    public static byte[] inetAddressTransform(byte[] classfileBuffer){
        ClassReader reader = new ClassReader(classfileBuffer);
        ClassNode node = new ClassNode(ASM5);
        reader.accept(node, ClassReader.EXPAND_FRAMES);
        for (MethodNode mn : node.methods) {
            if ("getAllByName".equals(mn.name) && "(Ljava/lang/String;Ljava/net/InetAddress;)[Ljava/net/InetAddress;".equals(mn.desc)) {
                InsnList list = new InsnList();
                list.add(new LabelNode());
                list.add(new VarInsnNode(ALOAD, 0));
                LabelNode label = new LabelNode();
                list.add(new JumpInsnNode(IFNONNULL, label));
                list.add(new InsnNode(ACONST_NULL));
                list.add(new InsnNode(ARETURN));
                list.add(label);
                list.add(new LabelNode());
                list.add(new LdcInsnNode("agent-addr"));
                list.add(new MethodInsnNode(INVOKESTATIC, "java/lang/System", "getProperty", "(Ljava/lang/String;)Ljava/lang/String;", false));
                list.add(new VarInsnNode(ASTORE, 3));
                list.add(new LabelNode());
                list.add(new VarInsnNode(ALOAD, 3));
                list.add(new VarInsnNode(ALOAD, 0));
                list.add(new MethodInsnNode(INVOKEVIRTUAL, "java/lang/String", "contains", "(Ljava/lang/CharSequence;)Z", false));
                LabelNode labelNode = new LabelNode();
                list.add(new JumpInsnNode(IFEQ, labelNode));
                list.add(new LabelNode());
                list.add(new TypeInsnNode(NEW, "java/net/UnknownHostException"));
                list.add(new InsnNode(DUP));
                list.add(new MethodInsnNode(INVOKESPECIAL, "java/net/UnknownHostException", "<init>", "()V", false));
                list.add(new InsnNode(ATHROW));
                list.add(labelNode);
                mn.instructions.insert(list);
                continue;
            }
            if ("isReachable".equals(mn.name) && "(Ljava/net/NetworkInterface;II)Z".equals(mn.desc)) {
                InsnList list = new InsnList();
                list.add(new LabelNode());
                list.add(new LdcInsnNode("agent-addr"));
                list.add(new MethodInsnNode(INVOKESTATIC, "java/lang/System", "getProperty", "(Ljava/lang/String;)Ljava/lang/String;", false));
                list.add(new VarInsnNode(ASTORE, 4));
                list.add(new LabelNode());
                list.add(new VarInsnNode(ALOAD, 4));
                list.add(new VarInsnNode(ALOAD, 0));
                list.add(new MethodInsnNode(INVOKEVIRTUAL, "java/net/InetAddress", "getHostName", "()Ljava/lang/String;", false));
                list.add(new MethodInsnNode(INVOKEVIRTUAL, "java/lang/String", "contains", "(Ljava/lang/CharSequence;)Z", false));
                LabelNode labelNode = new LabelNode();
                list.add(new JumpInsnNode(IFEQ, labelNode));
                list.add(new LabelNode());
                list.add(new InsnNode(ICONST_0));
                list.add(new InsnNode(IRETURN));
                list.add(labelNode);
                mn.instructions.insert(list);
            }
        }
        ClassWriter writer = new ClassWriter(ClassWriter.COMPUTE_FRAMES | ClassWriter.COMPUTE_MAXS);
        node.accept(writer);
        return writer.toByteArray();
    }
}
